package com.ruoyi.mqtt.service;

import com.ruoyi.mqtt.model.SensorData;
import lombok.extern.slf4j.Slf4j;
import org.apache.iotdb.isession.SessionDataSet;
import org.apache.iotdb.rpc.IoTDBConnectionException;
import org.apache.iotdb.rpc.StatementExecutionException;
import org.apache.iotdb.session.pool.SessionPool;
import org.apache.iotdb.tsfile.file.metadata.enums.TSDataType;
import org.apache.iotdb.tsfile.read.common.RowRecord;
import org.springframework.stereotype.Service;

import javax.annotation.PostConstruct;
import javax.annotation.PreDestroy;
import java.time.Instant;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicLong;

/**
 * IoTDB数据服务
 * 用于处理传感器数据的批量写入
 */
@Slf4j
@Service
public class IoTDBService {

    private SessionPool sessionPool;

    // 存储组路径前缀
    private static final String STORAGE_GROUP = "root.sensors";
    // 批处理大小
    private static final int BATCH_SIZE = 100;
    // 计数器，用于统计写入数量
    private final AtomicLong insertCount = new AtomicLong(0);
    // 重试次数
    private static final int MAX_RETRY_ATTEMPTS = 3;
    // 重试延迟（毫秒）
    private static final long RETRY_DELAY = 5000;

    // 添加一个去重的缓存
    private final ConcurrentHashMap<String, Long> lastInsertTime = new ConcurrentHashMap<>();

    @PostConstruct
    public void init() {
        try {
            // 使用SessionPool替代单个Session
            List<String> nodeUrls = new ArrayList<>();
            nodeUrls.add("127.0.0.1:6667");
            sessionPool = new SessionPool.Builder()
                    .nodeUrls(nodeUrls)
                    .user("root")
                    .password("root")
                    .maxSize(5) // 连接池大小
                    .build();

            initStorageGroup();
            log.info("IoTDB SessionPool初始化成功");
        } catch (Exception e) {
            log.error("Failed to initialize IoTDB SessionPool", e);
        }
    }

    @PreDestroy
    public void destroy() {
        if (sessionPool != null) {
            sessionPool.close();
            log.info("IoTDB SessionPool已关闭");
        }
    }

    /**
     * 初始化存储组
     * 在服务启动时创建IoTDB的存储组
     */
    public void initStorageGroup() {
        int retryAttempts = 0;
        while (retryAttempts < MAX_RETRY_ATTEMPTS) {
            try {
                sessionPool.setStorageGroup(STORAGE_GROUP);
                // 使用TTL设置
                String ttlSql = String.format("SET TTL TO %s %d", STORAGE_GROUP, 1000L * 60 * 60 * 24 * 7);
                sessionPool.executeNonQueryStatement(ttlSql);
                log.info("存储组 {} 初始化成功", STORAGE_GROUP);
                return;
            } catch (Exception e) {
                if (e instanceof org.apache.iotdb.rpc.StatementExecutionException
                    && e.getMessage().contains("has already been created as database")) {
                    log.info("存储组 {} 已存在，跳过初始化", STORAGE_GROUP);
                    return;
                }
                log.error("Failed to initialize storage group, attempt {}/{}",
                    retryAttempts + 1, MAX_RETRY_ATTEMPTS, e);
                retryAttempts++;
                if (retryAttempts < MAX_RETRY_ATTEMPTS) {
                    try {
                        Thread.sleep(RETRY_DELAY);
                    } catch (InterruptedException ie) {
                        Thread.currentThread().interrupt();
                        break;
                    }
                }
            }
        }
    }

    /**
     * 批量插入传感器数据
     *
     * @param dataList 传感器数据列表
     */
    public void batchInsertSensorData(List<SensorData> dataList) {
        if (dataList == null || dataList.isEmpty()) {
            return;
        }

        // 数据去重
        List<SensorData> uniqueDataList = new ArrayList<>();
        for (SensorData data : dataList) {
            String key = data.getSensorId() + "_" + data.getTimestamp().toEpochMilli();
            Long lastTime = lastInsertTime.get(key);
            if (lastTime == null || data.getTimestamp().toEpochMilli() > lastTime) {
                uniqueDataList.add(data);
                lastInsertTime.put(key, data.getTimestamp().toEpochMilli());
            }
        }

        if (uniqueDataList.isEmpty()) {
            return;
        }

        // 准备批量插入的数据结构
        List<String> deviceIds = new ArrayList<>();
        List<Long> timestamps = new ArrayList<>();
        List<List<String>> measurementsList = new ArrayList<>();
        List<List<TSDataType>> typesList = new ArrayList<>();
        List<List<Object>> valuesList = new ArrayList<>();

        int count = 0;
        int totalInserted = 0;

        for (SensorData data : uniqueDataList) {
            String deviceId = STORAGE_GROUP + "." + data.getSensorId().replace("-", "_");
            List<String> measurements = Arrays.asList("value", "roadName", "unit");
            List<TSDataType> types = Arrays.asList(
                TSDataType.DOUBLE,
                TSDataType.TEXT,
                TSDataType.TEXT
            );
            List<Object> values = Arrays.asList(
                data.getValue(),
                data.getRoadName(),
                data.getUnit()
            );

            deviceIds.add(deviceId);
            timestamps.add(data.getTimestamp().toEpochMilli());
            measurementsList.add(measurements);
            typesList.add(types);
            valuesList.add(values);

            count++;

            // 当达到批处理大小时执行插入
            if (count >= BATCH_SIZE) {
                try {
                    sessionPool.insertRecords(deviceIds, timestamps, measurementsList, typesList, valuesList);
                    totalInserted += count;
                    log.debug("成功插入{}条记录到IoTDB", count);
                } catch (IoTDBConnectionException | StatementExecutionException e) {
                    log.error("插入数据到IoTDB失败", e);
                }

                // 清空集合，准备下一批
                deviceIds.clear();
                timestamps.clear();
                measurementsList.clear();
                typesList.clear();
                valuesList.clear();
                count = 0;
            }
        }

        // 处理剩余的数据
        if (!deviceIds.isEmpty()) {
            try {
                sessionPool.insertRecords(deviceIds, timestamps, measurementsList, typesList, valuesList);
                totalInserted += count;
                log.debug("成功插入剩余{}条记录到IoTDB", count);
            } catch (IoTDBConnectionException | StatementExecutionException e) {
                log.error("插入剩余数据到IoTDB失败", e);
            }
        }

        // 强制刷新数据到磁盘
        try {
            sessionPool.executeNonQueryStatement("flush");
            long total = insertCount.addAndGet(totalInserted);
            log.info("成功插入{}条记录到IoTDB，总计: {}", totalInserted, total);
        } catch (IoTDBConnectionException | StatementExecutionException e) {
            log.error("刷新数据到磁盘失败", e);
        }
    }

    /**
     * 查询最新数据
     */
    public List<SensorData> queryLatestData(String sensorId, int limit) {
        List<SensorData> result = new ArrayList<>();
        try {
            String sql = String.format("SELECT * FROM root.sensors.%s ORDER BY time DESC LIMIT %d",
                sensorId.replace("-", "_"), limit);
            SessionDataSet sessionDataSet = sessionPool.executeQueryStatement(sql).getSessionDataSet();
            while (sessionDataSet.hasNext()) {
                RowRecord record = sessionDataSet.next();
                SensorData data = new SensorData();
                data.setSensorId(sensorId);
                data.setTimestamp(Instant.ofEpochMilli(record.getTimestamp()));

                // 设置其他字段
                if (record.getFields().size() >= 3) {
                    data.setValue(record.getFields().get(0).getDoubleV());
                    data.setRoadName(record.getFields().get(1).getStringValue());
                    data.setUnit(record.getFields().get(2).getStringValue());
                    result.add(data);
                }
            }
            sessionDataSet.closeOperationHandle();
        } catch (Exception e) {
            log.error("查询传感器{}的最新数据失败", sensorId, e);
        }
        return result;
    }

    /**
     * 获取数据总数
     */
    public long getDataCount(String sensorId) {
        try {
            String sql = String.format("SELECT COUNT(*) FROM root.sensors.%s",
                    sensorId.replace("-", "_"));
            SessionDataSet dataSet = sessionPool.executeQueryStatement(sql).getSessionDataSet();
            if (dataSet.hasNext()) {
                return dataSet.next().getFields().get(0).getLongV();
            }
        } catch (Exception e) {
            log.error("获取传感器{}的数据总数失败", sensorId, e);
        }
        return 0;
    }
}
